"""
框架3：成分数据（成分预测）
"""

import numpy as np
import pandas as pd
from typing import Dict, List, Tuple, Optional
import sys
import os

# 添加父目录到路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from ..data_loader import DataLoader
from ..ilr_transformation import ILRTransformation
from ..forecasting_models import ForecastingModels
from ..evaluation import Evaluation


class Framework3:
    """框架3：成分数据（成分预测）"""
    
    def __init__(self, data: Dict, config: Dict = None):
        """
        初始化框架3
        
        Args:
            data: 处理后的数据字典
            config: 配置字典
        """
        self.data = data
        self.config = config or {}
        self.forecasting_models = ForecastingModels(config)
        self.evaluation = Evaluation()
        
        # 提取数据
        self._extract_data()
        
    def _extract_data(self):
        """提取和准备数据"""
        # 获取第一个工作表的数据
        sheet_name = list(self.data.keys())[0]
        sheet_data = self.data[sheet_name]
        
        # 提取比例数据
        self.proportions_data = sheet_data['proportions']
        
        # 提取各省份在区域中的占比数据
        self.province_regional_prop = self.proportions_data['province_regional']
        
        # 提取区域在全国的占比数据
        self.regional_national_prop = self.proportions_data['regional_national']
        
    def prepare_compositional_data(self) -> pd.DataFrame:
        """
        准备成分数据
        
        Returns:
            成分数据矩阵
        """
        try:
            # 准备省份在区域中的占比数据
            province_regional_wide = self.province_regional_prop.pivot_table(
            index='Year',
                columns='Province',
                values='Regional_Proportion',
            fill_value=0
        )
        
            # 处理零值和负值问题
            epsilon = 1e-10
            province_regional_wide = province_regional_wide.where(province_regional_wide > 0, epsilon)
            
            # 按区域分组并重新归一化
            regions = {}
            # 使用硬编码的区域配置作为备用
            default_regions = {
                "华北地区": ["北京市", "河北+天津", "山西省", "内蒙古自治区"],
                "东北地区": ["辽宁省", "吉林省", "黑龙江省"],
                "华东地区": ["上海市", "江苏省", "浙江省", "安徽省", "福建省", "江西省", "山东省"],
                "华中地区": ["河南省", "湖北省", "湖南省"],
                "华南地区": ["广东+海南", "广西壮族自治区"],
                "西南地区": ["四川+重庆", "贵州省", "云南省", "西藏自治区"],
                "西北地区": ["陕西省", "甘肃省", "青海省", "宁夏回族自治区", "新疆维吾尔自治区"]
            }
            
            region_config = self.config.get('hierarchy', {}).get('regions', default_regions)
            
            for region_name, provinces in region_config.items():
                region_provinces = [p for p in provinces if p in province_regional_wide.columns]
                if region_provinces:
                    region_data = province_regional_wide[region_provinces]
                    # 重新归一化确保每行和为1
                    row_sums = region_data.sum(axis=1)
                    region_data = region_data.div(row_sums, axis=0)
                    regions[region_name] = region_data
                    print(f"    区域 {region_name}: {len(region_provinces)} 个省份")
            
            print(f"  处理了 {len(regions)} 个区域的数据")
            return regions
            
        except Exception as e:
            print(f"准备成分数据失败: {e}")
            # 返回简化的数据结构
            simplified_data = self.province_regional_prop.pivot_table(
                index='Year',
                columns='Province', 
                values='Regional_Proportion',
                fill_value=1e-10
            )
            return {'simplified': simplified_data}
    
    def forecast(self, horizon: int = 5) -> pd.DataFrame:
        """
        执行简化的ILR-VAR预测
        
        Args:
            horizon: 预测步长
            
        Returns:
            预测结果DataFrame
        """
        try:
            from statsmodels.tsa.arima.model import ARIMA
            import numpy as np
            
            print("Framework3: 开始ILR-VAR预测...")
            
        # 准备成分数据
            regions_data = self.prepare_compositional_data()
            
            if not regions_data:
                print("  没有可用的区域数据")
                return None
            
            # 获取所有省份的预测结果
            all_forecasts = {}
            
            # 对每个区域分别处理
            for region_name, region_data in regions_data.items():
                if region_name == 'simplified':
                    # 处理简化数据
                    print(f"  使用简化方法处理数据")
                    for province in region_data.columns:
                        series = region_data[province].dropna()
                        if len(series) >= 10:
                            try:
                                model = ARIMA(series, order=(1, 1, 1))
                                fitted_model = model.fit()
                                forecast_result = fitted_model.forecast(steps=horizon)
                                all_forecasts[province] = np.maximum(forecast_result.values, 1e-10)
                            except Exception as e:
                                print(f"    预测 {province} 失败: {str(e)}")
                                recent_mean = series.tail(3).mean()
                                all_forecasts[province] = np.full(horizon, max(recent_mean, 1e-10))
                    continue
                
                print(f"  处理区域: {region_name}")
                
                if region_data.shape[1] < 2:
                    print(f"    区域 {region_name} 省份数量不足，跳过")
                    continue
                
                # 对区域内的省份比例进行预测
                region_forecasts = {}
                
                print(f"    开始预测区域内 {len(region_data.columns)} 个省份的比例")
                
                for province in region_data.columns:
                    series = region_data[province].dropna()
                    
                    print(f"    处理 {province}: 数据点数={len(series)}, 最后值={series.iloc[-1] if len(series)>0 else 'N/A'}")
                    
                    if len(series) < 5:  # 降低要求
                        print(f"      跳过 {province}：数据点不足")
                        continue
                    
                    try:
                        # 使用ARIMA预测省份在区域中的比例
                        model = ARIMA(series, order=(1, 0, 0))  # 简化模型
                        fitted_model = model.fit()
                        forecast_result = fitted_model.forecast(steps=horizon)
                        
                        # 确保预测的比例为正数
                        forecast_values = np.maximum(forecast_result.values, 1e-10)
                        region_forecasts[province] = forecast_values
                        print(f"      ✓ ARIMA预测成功: {forecast_values}")
                        
                    except Exception as e:
                        print(f"      ARIMA预测 {province} 失败: {str(e)}")
                        # 使用最后几年的平均值作为预测
                        recent_values = series.tail(min(5, len(series)))
                        recent_mean = recent_values.mean()
                        if pd.isna(recent_mean) or recent_mean <= 0:
                            recent_mean = 1.0 / len(region_data.columns)  # 平均分配
                        
                        forecast_values = np.full(horizon, max(recent_mean, 1e-10))
                        region_forecasts[province] = forecast_values
                        print(f"      ✓ 使用均值预测: {forecast_values}")
                
                print(f"    成功预测了 {len(region_forecasts)} 个省份")
                
                if region_forecasts:
                    # 重新归一化区域内的比例
                    region_df = pd.DataFrame(region_forecasts)
                    print(f"    比例预测矩阵形状: {region_df.shape}")
                    print(f"    比例预测前5行:")
                    print(region_df.head())
                    
                    row_sums = region_df.sum(axis=1)
                    print(f"    行和: {row_sums.values}")
                    
                    if (row_sums > 0).all():
                        region_df = region_df.div(row_sums, axis=0)
                        print(f"    ✓ 归一化完成")
                    else:
                        print(f"    ⚠ 行和包含零值，使用平均分配")
                        region_df = region_df * 0 + (1.0 / len(region_df.columns))
                    
                    # 预测区域总量（简化处理）
                    try:
                        # 使用区域历史总量数据
                        regional_totals = self.regional_national_prop[
                            self.regional_national_prop['Area'] == region_name
                        ].set_index('Year')['Grain']
                        
                        print(f"    区域 {region_name} 历史数据点数: {len(regional_totals)}")
                        
                        if len(regional_totals) >= 10:
                            region_model = ARIMA(regional_totals, order=(1, 1, 1))
                            region_fitted = region_model.fit()
                            region_total_forecast = region_fitted.forecast(steps=horizon)
                            print(f"    区域总量预测: {region_total_forecast.values}")
                        else:
                            # 使用简单趋势外推
                            if len(regional_totals) > 0:
                                recent_growth = regional_totals.tail(min(5, len(regional_totals))).pct_change().mean()
                                if pd.isna(recent_growth):
                                    recent_growth = 0.02  # 默认2%增长
                                last_value = regional_totals.iloc[-1]
                                region_total_forecast = [last_value * (1 + recent_growth) ** (i+1) for i in range(horizon)]
                                region_total_forecast = pd.Series(region_total_forecast)
                                print(f"    使用趋势外推，增长率: {recent_growth:.3f}, 预测: {region_total_forecast.values}")
                            else:
                                # 没有历史数据，使用全国数据估算
                                print(f"    区域 {region_name} 无历史数据，使用估算")
                                estimated_total = 10000  # 估算值
                                region_total_forecast = pd.Series([estimated_total * (1.02 ** (i+1)) for i in range(horizon)])
                        
                        # 将区域内比例转换为绝对值
                        print(f"    开始转换比例为绝对值...")
                        print(f"    region_df形状: {region_df.shape}")
                        print(f"    region_total_forecast形状: {region_total_forecast.shape}")
                        
                        for province in region_df.columns:
                            # 确保索引长度匹配
                            province_proportions = region_df[province].values
                            region_totals = region_total_forecast.values
                            
                            # 计算绝对值
                            province_forecast = province_proportions * region_totals
                            
                            # 确保没有NaN值
                            province_values = np.where(np.isnan(province_forecast), 0, province_forecast)
                            all_forecasts[province] = province_values
                            print(f"      {province}: 比例={province_proportions} × 总量={region_totals} = {province_values}")
                            
                    except Exception as e:
                        print(f"    区域 {region_name} 总量预测失败: {str(e)}")
                        # 使用更稳健的备用方案
                        for province in region_df.columns:
                            # 基于省份历史数据的简单预测
                            province_data = region_data[province].dropna()
                            if len(province_data) > 0:
                                last_value = province_data.iloc[-1]
                                # 转换比例为绝对值（假设区域总量约10000万吨）
                                estimated_values = [last_value * 10000 * (1.02 ** (i+1)) for i in range(horizon)]
                                all_forecasts[province] = np.array(estimated_values)
                                print(f"      {province} (备用方案): {estimated_values}")
            
            if not all_forecasts:
                print("  没有成功预测任何省份")
                return None
            
            # 构建最终预测结果
            result_df = pd.DataFrame(all_forecasts)
            
            # 设置正确的时间索引
            if 'simplified' in regions_data:
                last_year = regions_data['simplified'].index.max()
            else:
                # 从province_regional_prop获取最后年份
                last_year = self.province_regional_prop['Year'].max()
            
            # 确保索引长度匹配
            if len(result_df) > 0:
                actual_horizon = len(result_df)
                future_years = list(range(last_year + 1, last_year + actual_horizon + 1))
                result_df.index = future_years
            
            print(f"  Framework3预测完成，结果形状: {result_df.shape}")
            return result_df
            
        except Exception as e:
            print(f"Framework3预测失败: {str(e)}")
            import traceback
            traceback.print_exc()
            return None
    
    def _reconstruct_forecasts(self, forecast_results: pd.DataFrame, horizon: int) -> Dict:
        """
        重建预测结果
        
        Args:
            forecast_results: 预测结果数据框
            horizon: 预测步长
            
        Returns:
            重建后的预测结果字典
        """
        # 分离区域预测和ILR预测
        regional_columns = [col for col in forecast_results.columns if col in self.regional_national_prop['Area'].unique()]
        ilr_columns = [col for col in forecast_results.columns if 'ILR' in col]
        
        regional_forecasts = forecast_results[regional_columns]
        ilr_forecasts = forecast_results[ilr_columns]
        
        # 重建省份预测
        province_forecasts = self._reconstruct_province_forecasts(
            regional_forecasts, ilr_forecasts, horizon
        )
        
        # 计算全国总量预测
        national_forecast = regional_forecasts.sum(axis=1)
        
        return {
            'national': national_forecast,
            'regional': regional_forecasts,
            'province': province_forecasts
        }
    
    def _reconstruct_province_forecasts(self, regional_forecasts: pd.DataFrame, 
                                     ilr_forecasts: pd.DataFrame, horizon: int) -> pd.DataFrame:
        """
        重建省份预测
        
        Args:
            regional_forecasts: 区域预测结果
            ilr_forecasts: ILR预测结果
            horizon: 预测步长
            
        Returns:
            省份预测结果
        """
        # 这里需要根据ILR逆变换重建各省份的预测
        # 由于ILR变换的复杂性，这里提供一个简化的实现
        
        province_forecasts = {}
        
        # 按区域重建省份预测
        for region in regional_forecasts.columns:
            region_forecast = regional_forecasts[region].values
            
            # 获取该区域的省份列表
            region_provinces = self.province_regional_prop[
                self.province_regional_prop['Area'] == region
            ]['Province'].unique()
            
            # 简化的重建方法：使用历史平均比例
            historical_proportions = self.province_regional_prop[
                self.province_regional_prop['Area'] == region
            ].groupby('Province')['Regional_Proportion'].mean()
            
            # 重建各省份预测
            for province in region_provinces:
                if province in historical_proportions.index:
                    proportion = historical_proportions[province]
                    province_forecasts[f"{province}_{region}"] = region_forecast * proportion
        
        return pd.DataFrame(province_forecasts)
    
    def run_complete_analysis(self, horizons: List[int] = None) -> Dict:
        """
        运行完整的分析流程
        
        Args:
            horizons: 预测步长列表
            
        Returns:
            完整分析结果
        """
        if horizons is None:
            horizons = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 12, 15, 20]
        
        all_results = {}
        
        for horizon in horizons:
            print(f"正在处理预测步长 h={horizon}")
            
            # 执行预测
            forecast_results = self.forecast(horizon=horizon)
            
            all_results[f"h={horizon}"] = {
                'forecast': forecast_results
            }
        
        return all_results 